from ultralytics import YOLO
import os

# 定义模型路径
model_path = 'yolov8n.pt'
# 加载预训练的YOLOv8n模型
model = YOLO(model_path)

# 输入图片所在的文件夹
input_folder = '照片/srgan后的'
# ，如果不存在则创建
output_folder = 'output_results'
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 遍历input_folder中的所有文件
for filename in os.listdir(input_folder):
    if filename.lower().endswith(('.png', '.jpg', '.jpeg')):  # 检查是否为图片文件
        # 构建图片文件的完整路径
        img_path = os.path.join(input_folder, filename)

        # 使用模型进行预测
        results = model(img_path, stream=False)  # 对于单张图片，stream=False

        # 检查结果列表是否不为空且包含Results对象
        if results:
            result = results[0]  # 从列表中取出第一个（也是唯一的）Results对象

            # 构建输出图片的保存路径
            output_img_path = os.path.join(output_folder, filename)

            # 保存结果到文件
            result.save(filename=output_img_path)

print('所有图片的处理和保存已完成。')